#!/usr/bin/python
# copy from respirenet
# original author: Siddahartha Gairola (t-sigai at microsoft dot com)

import numpy as np
import os
import io
import math
import random
import pandas as pd

import matplotlib.pyplot as plt
import librosa
import librosa.display
import cv2
import cmapy

import nlpaug
import nlpaug.augmenter.audio as naa

import torch

from scipy.signal import butter, lfilter

def Extract_Annotation_Data(file_name, data_dir):
    tokens = file_name.split('_')
    recording_info = pd.DataFrame(
        data = [tokens], 
        columns = [
            'Patient Number', 
            'Recording index', 
            'Chest location',
            'Acquisition mode',
            'Recording equipment'
        ]
    )
    recording_annotations = pd.read_csv(
        os.path.join(data_dir, file_name + '.txt'), 
        names = ['Start', 'End', 'Crackles', 'Wheezes'], 
        delimiter= '\t'
    )
    return recording_info, recording_annotations

#Used to split each individual sound file into separate sound clips containing one respiratory cycle each
#output: [filename, (sample_data:np.array, start:float, end:float, label:int (...) ]
#label: 0-normal, 1-crackle, 2-wheeze, 3-both
def get_sound_samples(recording_annotations, file_name, data_dir, sample_rate):
    sample_data = [file_name]
    # load file with specified sample rate (also converts to mono)
    data, rate = librosa.load(os.path.join(data_dir, file_name+'.wav'), sr=sample_rate)
    #print("Sample Rate", rate)
    
    for i in range(len(recording_annotations.index)):
        row = recording_annotations.loc[i]
        start = row['Start']
        end = row['End']
        crackles = row['Crackles']
        wheezes = row['Wheezes']
        audio_chunk = slice_data(start, end, data, rate)
        sample_data.append((audio_chunk, start,end, get_label(crackles, wheezes)))
    return sample_data


# split samples according to desired length
'''
types:
* 0: simply pad by zeros
* 1: pad with duplicate on both sides (half-n-half)
* 2: pad with augmented sample on both sides (half-n-half)	
'''
def Split_and_Pad(original, desiredLength, sample_rate, types=0):
	if types==0:
		return split_and_pad_old(original, desiredLength, sample_rate)

	output_buffer_length = int(desiredLength*sample_rate)
	soundclip = original[0].copy()
	n_samples = len(soundclip)

	output = []
	# if: the audio sample length > desiredLength, then split & pad
	# else: simply pad according to given type 1 or 2
	if n_samples > output_buffer_length:
		frames = librosa.util.frame(soundclip, frame_length=output_buffer_length, hop_length=output_buffer_length//2, axis=0)
		for i in range(frames.shape[0]):
			output.append((frames[i], original[1], original[2], original[3], original[4], i, 0))

		last_id = frames.shape[0]*(output_buffer_length//2)
		last_sample = soundclip[last_id:]; pad_times = (output_buffer_length-len(last_sample))/len(last_sample)
		padded = generate_padded_samples(soundclip, last_sample, output_buffer_length, sample_rate, types)
		output.append((padded, original[1], original[2], original[3], original[4], i+1, pad_times))

	else:
		padded = generate_padded_samples(soundclip, soundclip, output_buffer_length, sample_rate, types); pad_times = (output_buffer_length-len(soundclip))/len(soundclip)
		output.append((padded, original[1], original[2], original[3], original[4], 0, pad_times))

	return output


def split_and_pad_old(original, desiredLength, sample_rate):
    output_buffer_length = int(desiredLength * sample_rate)
    soundclip = original[0].copy()
    n_samples = len(soundclip)
    total_length = n_samples / sample_rate #length of cycle in seconds
    n_slices = int(math.ceil(total_length / desiredLength)) #get the minimum number of slices needed
    samples_per_slice = n_samples // n_slices
    src_start = 0 #Staring index of the samples to copy from the original buffer
    output = [] #Holds the resultant slices
    for i in range(n_slices):
        src_end = min(src_start + samples_per_slice, n_samples)
        length = src_end - src_start
        copy = generate_padded_samples_old(soundclip[src_start:src_end], output_buffer_length)
        output.append((copy, original[1], original[2], original[3], original[4], i))
        src_start += length
    return output


def generate_padded_samples_old(source, output_length):
    copy = np.zeros(output_length, dtype = np.float32)
    src_length = len(source)
    frac = src_length / output_length
    if(frac < 0.5):
        #tile forward sounds to fill empty space
        cursor = 0
        while(cursor + src_length) < output_length:
            copy[cursor:(cursor + src_length)] = source[:]
            cursor += src_length
    else:
        copy[:src_length] = source[:]
    return copy

def generate_padded_samples(original, source, output_length, sample_rate, types):
	copy = np.zeros(output_length, dtype=np.float32)
	src_length = len(source)
	left = output_length-src_length # amount to be padded
	# pad front or back
	prob = random.random()
	if types == 1:
		aug = original
	else:
		aug = gen_augmented(original, sample_rate)

	while len(aug) < left:
		aug = np.concatenate([aug, aug])

	if prob < 0.5:
		#pad back
		copy[left:] = source
		copy[:left] = aug[len(aug)-left:]
	else:
		#pad front
		copy[:src_length] = source[:]
		copy[src_length:] = aug[:left]

	return copy

#**********************DATA AUGMENTAION***************************
#Creates a copy of each time slice, but stretches or contracts it by a random amount
def gen_augmented(original, sample_rate):
	# list of augmentors available from the nlpaug library
	augment_list = [
	#naa.CropAug(sampling_rate=sample_rate)
	naa.NoiseAug(),
	naa.SpeedAug(),
	naa.LoudnessAug(factor=(0.5, 2)),
	naa.VtlpAug(sampling_rate=sample_rate, zone=(0.0, 1.0)),
	naa.PitchAug(sampling_rate=sample_rate, factor=(-1,3))
	]
	# sample augmentation randomly
	aug_idx = random.randint(0, len(augment_list)-1)
	augmented_data = augment_list[aug_idx].augment(original)
	return augmented_data
